# 导入必要的包
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
from tensorflow.keras.models import Sequential, load_model, save_model
from tensorflow.keras.layers import LSTM, Dense
from sklearn.preprocessing import MinMaxScaler
from sklearn.linear_model import LogisticRegression
import pickle
import os
from datetime import datetime
import matplotlib.dates as mdates

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


class RainfallForecaster:
    """常州降水预测系统"""

    def __init__(self, data_file='常州.csv', time_step=10, epochs=20, batch_size=32):
        """
        初始化降水预测系统
        """
        self.data_file = data_file
        self.time_step = time_step
        self.epochs = epochs
        self.batch_size = batch_size

        # 提取文件名作为基础名称（不含路径和扩展名）
        self.base_filename = os.path.splitext(os.path.basename(data_file))[0]

        # 模型文件路径
        self.lstm_model_path = f'{self.base_filename}_lstm_model.h5'
        self.lr_model_path = f'{self.base_filename}_lr_model.pkl'
        self.scaler_path = f'{self.base_filename}_scaler.pkl'
        self.info_path = f'{self.base_filename}_model_info.pkl'

        # 初始化变量
        self.df = None
        self.model_lstm = None
        self.model_lr = None
        self.scaler = None
        self.model_info = None

        # 修改为用户提供的列名映射
        self.column_mapping = {
            'temperture': 'temperature',  # 保持原拼写以匹配输入数据
            '气压': 'pressure',
            '降雨': 'rainfall',
            '风速': 'wind_speed',
            'date': 'date'  # 日期列
        }

        print(f"初始化常州降水预测系统")
        print(f"数据文件: {data_file}")
        print(f"模型文件将保存为: {self.lstm_model_path}")

    def load_data(self):
        """加载常州气象数据"""
        try:
            # 读取CSV文件
            df = pd.read_csv(self.data_file)
            # 修正列名（去除空格）
            df.columns = [col.strip() for col in df.columns]

            print(f"原始列名: {list(df.columns)}")

            # 确保日期列存在
            if 'date' not in df.columns:
                raise ValueError("数据文件缺少'date'列")

            # 将日期列转换为datetime类型并设为索引
            df['date'] = pd.to_datetime(df['date'])
            df.set_index('date', inplace=True)
            # 按时间排序
            df.sort_index(inplace=True)

            # 重命名列以匹配预期名称
            for orig_col, new_col in self.column_mapping.items():
                if orig_col in df.columns:
                    df[new_col] = df[orig_col]
                else:
                    print(f"警告: 找不到'{orig_col}'列，创建随机数据")
                    if new_col == 'rainfall':
                        df[new_col] = np.random.gamma(1, 5, len(df))
                    elif new_col == 'temperature':
                        df[new_col] = np.random.normal(25, 5, len(df))
                    elif new_col == 'pressure':
                        df[new_col] = np.random.normal(1013, 10, len(df))
                    elif new_col == 'wind_speed':
                        df[new_col] = np.random.normal(15, 5, len(df))

            # 确保湿度数据可用（为兼容原有代码）
            if 'humidity' not in df.columns:
                print("警告: 找不到湿度数据，创建替代数据")
                df['humidity'] = np.random.normal(60, 10, len(df))

            self.df = df
            print(f"成功加载常州数据，共 {len(df)} 条记录")
        except Exception as e:
            print(f"加载常州数据出错: {e}")
            print("使用生成的示例数据代替")
            self.df = self._generate_example_data()

        return self.df

    def _generate_example_data(self):
        """生成示例气象数据（内部方法）"""
        date_rng = pd.date_range(start='2023-01-01', end='2023-12-31', freq='D')
        rainfall = np.random.gamma(1, 5, len(date_rng))
        temperature = np.random.normal(25, 5, len(date_rng))
        humidity = np.random.normal(60, 10, len(date_rng))
        pressure = np.random.normal(1013, 10, len(date_rng))
        wind_speed = np.random.normal(15, 5, len(date_rng))

        df = pd.DataFrame({
            'date': date_rng,
            'rainfall': rainfall,
            'temperature': temperature,
            'humidity': humidity,
            'pressure': pressure,
            'wind_speed': wind_speed
        })
        df.set_index('date', inplace=True)
        print("已生成示例数据")
        return df

    def _create_dataset(self, dataset):
        """创建时间序列数据集（内部方法）"""
        X, Y = [], []
        for i in range(len(dataset) - self.time_step - 1):
            a = dataset[i:(i + self.time_step), 0]
            X.append(a)
            Y.append(dataset[i + self.time_step, 0])
        return np.array(X), np.array(Y)

    def train_models(self):
        """训练降水预测模型"""
        print("开始训练降水预测模型...")

        # 确保数据已加载
        if self.df is None:
            self.load_data()

        # 确保数据中有必要的列
        required_columns = ['rainfall', 'temperature', 'pressure', 'wind_speed']
        for col in required_columns:
            if col not in self.df.columns:
                print(f"警告: 数据中缺少'{col}'列，创建替代数据")
                if col == 'rainfall':
                    self.df[col] = np.random.gamma(1, 5, len(self.df))
                elif col == 'temperature':
                    self.df[col] = np.random.normal(25, 5, len(self.df))
                elif col == 'pressure':
                    self.df[col] = np.random.normal(1013, 10, len(self.df))
                elif col == 'wind_speed':
                    self.df[col] = np.random.normal(15, 5, len(self.df))

        # 确保数据中没有缺失值
        self.df = self.df.dropna(subset=required_columns)

        # 降水概率预测（逻辑回归模型）- 使用温度、气压和风速
        X = self.df[['temperature', 'pressure', 'wind_speed']]
        y = (self.df['rainfall'] > 0).astype(int)  # 有降水为1，无降水为0

        self.model_lr = LogisticRegression()
        self.model_lr.fit(X, y)
        self.df['rainfall_prob'] = self.model_lr.predict_proba(X)[:, 1]
        print("逻辑回归模型训练完成")

        # 数据预处理
        self.scaler = MinMaxScaler(feature_range=(0, 1))
        scaled_data = self.scaler.fit_transform(self.df[['rainfall']])

        # 创建训练数据
        X_lstm, Y_lstm = self._create_dataset(scaled_data)
        X_lstm = X_lstm.reshape(X_lstm.shape[0], X_lstm.shape[1], 1)

        # 构建LSTM模型
        self.model_lstm = Sequential()
        self.model_lstm.add(LSTM(50, return_sequences=True, input_shape=(self.time_step, 1)))
        self.model_lstm.add(LSTM(50))
        self.model_lstm.add(Dense(1))
        self.model_lstm.compile(loss='mean_squared_error', optimizer='adam')

        # 训练模型
        history = self.model_lstm.fit(X_lstm, Y_lstm, epochs=self.epochs,
                                      batch_size=self.batch_size, verbose=1)
        print(f"LSTM模型训练完成，共训练 {self.epochs} 轮")

        # 保存模型信息
        self.model_info = {
            'time_step': self.time_step,
            'last_date': self.df.index[-1],
            'trained_date': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }

        # 保存模型
        self.save_models()

        return history

    def save_models(self):
        """保存训练好的模型和相关信息"""
        # 保存LSTM模型
        save_model(self.model_lstm, self.lstm_model_path)
        print(f"LSTM模型已保存到 {self.lstm_model_path}")

        # 保存逻辑回归模型
        with open(self.lr_model_path, 'wb') as f:
            pickle.dump(self.model_lr, f)
        print(f"逻辑回归模型已保存到 {self.lr_model_path}")

        # 保存scaler
        with open(self.scaler_path, 'wb') as f:
            pickle.dump(self.scaler, f)
        print(f"数据标准化器已保存到 {self.scaler_path}")

        # 保存模型信息
        with open(self.info_path, 'wb') as f:
            pickle.dump(self.model_info, f)
        print(f"模型信息已保存到 {self.info_path}")

    def load_models(self):
        """加载保存的模型"""
        try:
            # 加载LSTM模型
            self.model_lstm = load_model(self.lstm_model_path)

            # 加载逻辑回归模型
            with open(self.lr_model_path, 'rb') as f:
                self.model_lr = pickle.load(f)

            # 加载scaler
            with open(self.scaler_path, 'rb') as f:
                self.scaler = pickle.load(f)

            # 加载模型信息
            with open(self.info_path, 'rb') as f:
                self.model_info = pickle.load(f)

            print("模型加载成功")
            print(f"模型训练日期: {self.model_info.get('trained_date', '未知')}")
            return True
        except Exception as e:
            print(f"加载模型时出错: {e}")
            return False

    def predict(self, days=7):
        """预测未来降水量和降水概率"""
        # 检查是否需要加载数据
        if self.df is None:
            self.load_data()

        # 检查是否需要加载或训练模型
        if self.model_lstm is None or self.model_lr is None or self.scaler is None:
            if os.path.exists(self.lstm_model_path) and self.load_models():
                print("使用已保存的模型进行预测")
            else:
                print("未找到模型或加载失败，训练新模型...")
                self.train_models()

        # 提取时间步长
        time_step = self.model_info.get('time_step', self.time_step) if self.model_info else self.time_step

        # 准备预测数据
        scaled_data = self.scaler.transform(self.df[['rainfall']])
        last_sequence = scaled_data[-time_step:]
        last_sequence = last_sequence.reshape(1, time_step, 1)

        # 预测未来几天
        future_dates = []
        future_predictions = []
        future_probabilities = []

        current_sequence = last_sequence.copy()

        print(f"开始预测未来 {days} 天降水情况...")
        for i in range(days):
            # 预测下一天降水量
            next_rainfall = self.model_lstm.predict(current_sequence, verbose=0)[0][0]
            future_predictions.append(next_rainfall)

            # 更新序列用于下一次预测
            temp_seq = current_sequence[0].copy()
            temp_seq = np.roll(temp_seq, -1, axis=0)
            temp_seq[-1] = next_rainfall
            current_sequence = temp_seq.reshape(1, time_step, 1)

            # 生成下一天日期
            next_date = self.df.index[-1] + pd.Timedelta(days=i + 1)
            future_dates.append(next_date)

            # 使用最近的天气数据来预测降水概率
            temp = self.df['temperature'].iloc[-1]
            press = self.df['pressure'].iloc[-1]
            wind = self.df['wind_speed'].iloc[-1]

            # 添加一些随机变化，模拟天气变化
            temp += np.random.normal(0, 2)
            press += np.random.normal(0, 3)
            wind += np.random.normal(0, 1)

            # 预测降水概率
            rain_prob = self.model_lr.predict_proba(np.array([[temp, press, wind]]))[0][1]
            future_probabilities.append(rain_prob)

        # 反向转换预测结果
        future_predictions_array = np.array(future_predictions).reshape(-1, 1)
        future_rainfall = self.scaler.inverse_transform(future_predictions_array)

        # 创建未来预测数据框
        future_df = pd.DataFrame({
            'date': future_dates,
            'predicted_rainfall': future_rainfall.flatten(),
            'rainfall_probability': future_probabilities
        })
        future_df.set_index('date', inplace=True)

        # 保存预测结果
        csv_path = f'{self.base_filename}.csv'
        future_df.to_csv(csv_path)
        print(f"预测结果已保存到 {csv_path}")

        return future_df

    def visualize_results(self, future_df, history_days=15):
        """可视化预测结果"""
        if future_df is None or self.df is None:
            print("没有可视化的数据")
            return

        days = len(future_df)

        # 增强版可视化 - 使用子图和更好的格式
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 12), sharex=True)

        # 子图1：降水量
        # 历史降水量
        ax1.bar(self.df.index[-history_days:], self.df['rainfall'][-history_days:],
                width=0.7, color='royalblue', alpha=0.7, label='历史降水量')

        # 预测降水量
        prediction_bars = ax1.bar(future_df.index, future_df['predicted_rainfall'],
                                  width=0.7, color='crimson', alpha=0.7, label='预测降水量')

        # 添加数据标签
        for i, bar in enumerate(prediction_bars):
            height = bar.get_height()
            if height > 0.1:  # 只有当降水量大于0.1时才显示标签
                ax1.text(bar.get_x() + bar.get_width() / 2., height + 0.1,
                         f'{height:.1f}', ha='center', va='bottom', fontsize=9)

        # 设置子图1的标题和标签
        ax1.set_title('常州市降水量预测', fontsize=16, pad=10)
        ax1.set_ylabel('降水量 (mm)', fontsize=12)
        ax1.grid(True, linestyle='--', alpha=0.7)
        ax1.legend(loc='upper left')

        # 子图2：降水概率
        # 历史降水概率
        if 'rainfall_prob' in self.df.columns:
            ax2.plot(self.df.index[-history_days:], self.df['rainfall_prob'][-history_days:] * 100,
                     marker='o', linestyle='-', color='forestgreen', alpha=0.8,
                     label='历史降水概率')

        # 预测降水概率
        probs = future_df['rainfall_probability'] * 100
        prob_line = ax2.plot(future_df.index, probs,
                             marker='D', linestyle='-', linewidth=2,
                             color='orangered', alpha=0.8, label='预测降水概率')

        # 添加概率数据标签
        for i, (date, prob) in enumerate(zip(future_df.index, probs)):
            ax2.text(date, prob + 2, f'{prob:.1f}%',
                     ha='center', va='bottom', color='black', fontsize=9)

        # 用颜色区分不同概率区域
        ax2.axhspan(0, 30, alpha=0.2, color='lightgreen', label='低概率区')
        ax2.axhspan(30, 70, alpha=0.2, color='khaki', label='中概率区')
        ax2.axhspan(70, 100, alpha=0.2, color='salmon', label='高概率区')

        # 设置子图2的标题和标签
        ax2.set_title('降水概率预测', fontsize=16, pad=10)
        ax2.set_ylabel('降水概率 (%)', fontsize=12)
        ax2.set_ylim(0, 100)
        ax2.grid(True, linestyle='--', alpha=0.7)
        ax2.legend(loc='upper left')

        # 设置共享的x轴
        fig.autofmt_xdate()  # 自动格式化x轴日期标签
        date_format = mdates.DateFormatter('%m-%d')
        ax2.xaxis.set_major_formatter(date_format)
        ax2.xaxis.set_major_locator(mdates.DayLocator(interval=1))
        ax2.set_xlabel('日期', fontsize=12)

        # 整体布局
        fig.suptitle(f'常州市未来{days}天降水预测', fontsize=18, y=0.98)
        plt.tight_layout()
        fig.subplots_adjust(top=0.9)

        # 保存图表
        plt.savefig(f'{self.base_filename}_prediction.png', dpi=300, bbox_inches='tight')
        plt.show()

    def run(self, days=7, force_train=False):
        """运行完整的预测流程"""
        # 加载数据
        self.load_data()

        # 训练或加载模型
        if force_train or not os.path.exists(self.lstm_model_path) or not self.load_models():
            self.train_models()

        # 预测未来降水
        future_df = self.predict(days=days)

        # 可视化结果
        self.visualize_results(future_df)

        return future_df


# 主函数示例
if __name__ == "__main__":
    # 创建预测器实例
    forecaster = RainfallForecaster(data_file='./merged_output/7.csv')

    # 运行完整流程
    predictions = forecaster.run(days=7)

    print("\n未来降水预测结果:")
    print(predictions)